package edu.wisc.hevea.broker;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentSkipListSet;

import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import edu.wisc.cloudservice.CloudServiceClient;
import edu.wisc.hevea.constants.ApplicationConstants;
import edu.wisc.hevea.util.CommonUtil;
import edu.wisc.hevea.zookeeper.util.ZooKeeperUtil;

public class HeartBeatThread implements Runnable, Watcher {

	private static final Logger logger = LoggerFactory
			.getLogger(HeartBeatThread.class);

	private MessageBroker messageBroker;
	private String rootNode;
	private ZooKeeper currentZooKeeper;

	private Map<String, Date> dcScaledTimeMap;
	private Date brokerBootTime;

	public HeartBeatThread(MessageBroker messageBroker) {
		this.messageBroker = messageBroker;
		this.rootNode = messageBroker.getApplicationPolicy().getApplicationId();
		this.dcScaledTimeMap = new HashMap<String, Date>();
		brokerBootTime = new Date();
	}

	@Override
	public void run() {

		int sessionTimeout = 3000;
		String myHostName = messageBroker.getMyHostName();
		String zooKeeperhostname = ZooKeeperUtil
				.getPreferredZookeeperHostName(myHostName);

		try {

			currentZooKeeper = ZooKeeperUtil.getZooKeeperInstance(
					zooKeeperhostname, sessionTimeout, this);

			if (!ZooKeeperUtil.doesExistPath(currentZooKeeper,
					getAllBrokersPath())) {
				ZooKeeperUtil.createPath(currentZooKeeper, getAllBrokersPath(),
						ApplicationConstants.ALL_BROKERS);
			}
			if (!ZooKeeperUtil.doesExistPath(currentZooKeeper,
					getAllBrokersPath() + "/" + myHostName)) {
				ZooKeeperUtil.createPath(currentZooKeeper, getAllBrokersPath()
						+ "/" + myHostName, myHostName);
			}

			if (!ZooKeeperUtil.doesExistPath(currentZooKeeper,
					getAliveBrokersPath())) {
				ZooKeeperUtil.createPath(currentZooKeeper,
						getAliveBrokersPath(),
						ApplicationConstants.ALIVE_BROKERS);
			}
			if (!ZooKeeperUtil.doesExistPath(currentZooKeeper,
					getAliveBrokersPath() + "/" + myHostName)) {
				
				ZooKeeperUtil.createEphemeralPath(currentZooKeeper,
						getAliveBrokersPath() + "/" + myHostName, myHostName);
			}

			// To set watcher on the zookeeper paths
			/*
			 * ZooKeeperUtil.getChildren( currentZooKeeper, getAllBrokersPath(),
			 * this); ZooKeeperUtil.getChildren(currentZooKeeper,
			 * getAliveBrokersPath(), this);
			 */

			handleAliveBrokerChanged();

		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (KeeperException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (InterruptedException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}

	}

	private String getAllBrokersPath() {
		return "/" + rootNode + "/" + ApplicationConstants.ALL_BROKERS;
	}

	private String getAliveBrokersPath() {
		return "/" + rootNode + "/" + ApplicationConstants.ALIVE_BROKERS;
	}

	@Override
	public void process(WatchedEvent event) {
		if (event.getPath() == null) {
			return;
		}
		if (event.getPath().equals(getAliveBrokersPath())) {
			handleAliveBrokerChanged();
		}
	}

	private String getRandomBroker(List<String> aliveBrokers) {
		if (aliveBrokers.size() == 0) {
			return "";
		}
		int index = (int) (Math.random() * aliveBrokers.size());
		logger.info("Returning random broker : " + aliveBrokers.get(index));
		return aliveBrokers.get(index);
	}

	private synchronized void handleAliveBrokerChanged() {
		try {
			List<String> allBrokers = ZooKeeperUtil.getChildren(
					currentZooKeeper, getAllBrokersPath(), this);
			List<String> aliveBrokers = ZooKeeperUtil.getChildren(
					currentZooKeeper, getAliveBrokersPath(), this);

			Set<String> aliveBrokersSet = new HashSet<String>();
			aliveBrokersSet.addAll(aliveBrokers);
			messageBroker.setAliveBrokers(aliveBrokersSet);
			
			Date currentTime = new Date();

			int numBrokersToReplicate = messageBroker.getApplicationPolicy()
					.getReplication() - 1;

			ConcurrentSkipListSet<String> repBrokers = messageBroker
					.getReplicationBrokers();
			
			int numUniqueBrokersInSameDC = 0;
			Set<String> uniqueBrokersInSameDC = new HashSet<String>();
			for(String currentPreferredBroker : repBrokers) {
				if(ApplicationConstants.HOST_DC_MAPPING.get(currentPreferredBroker)
						.equals(ApplicationConstants.HOST_DC_MAPPING.get(messageBroker.getMyHostName())) &&
						!currentPreferredBroker.equals(messageBroker.getMyHostName())) {
					uniqueBrokersInSameDC.add(currentPreferredBroker);
				}
				if(!aliveBrokers.contains(currentPreferredBroker)) {
					messageBroker.getMessagingServiceClient().removeServiceInterface(currentPreferredBroker);
				}
			}
			numUniqueBrokersInSameDC = uniqueBrokersInSameDC.size();
			
			if (aliveBrokers.containsAll(repBrokers)
					&& numUniqueBrokersInSameDC == numBrokersToReplicate) {
				logger.info("NO_NEED_TO_CHANGE : hevea###----List of Replication brokers-------");
				for (String broker : repBrokers) {
					logger.info(" Broker : " + broker);
				}
				logger.info("hevea###----End Of Replication brokers-------");
				// No need to change the set of replicated brokers
			} else {
				String myDataCenter = ApplicationConstants.HOST_DC_MAPPING
						.get(messageBroker.getMyHostName());

				repBrokers.clear();
				List<String> aliveBrokersExceptMyself = new ArrayList<String>(aliveBrokers.size()-1);
				for(String broker: aliveBrokers) {
					if(!broker.equals(messageBroker.getMyHostName())) {
						aliveBrokersExceptMyself.add(broker);
					}
				}
				
				for (String broker : aliveBrokersExceptMyself) {
					if (ApplicationConstants.HOST_DC_MAPPING.get(broker)
							.equals(myDataCenter)
							&& repBrokers.size() < numBrokersToReplicate) {
						repBrokers.add(broker);
					}
				}

				if (repBrokers.size() < numBrokersToReplicate
						&& aliveBrokersExceptMyself.size() > 0) {
					while (repBrokers.size() < numBrokersToReplicate) {
						repBrokers
								.add(getRandomBroker(aliveBrokersExceptMyself));
					}
				}
				
				logger.info("hevea###----List of Replication brokers-------");
				for (String broker : repBrokers) {
					logger.info(" Broker : " + broker);
				}
				logger.info("hevea###----End Of Replication brokers-------");
			}

			if (!messageBroker.isLeader()
					|| !CommonUtil.canInitiateScalingAfterBoot(brokerBootTime,
							currentTime)) {
				return;
			}

			logger.info("hevea###----List of Alivebrokers-------");
			for (String broker : aliveBrokers) {
				logger.info(" Broker : " + broker);				
			}
			logger.info("hevea###----End of Alivebrokers-------");

			int numBrokersRequiredPerDC = messageBroker.getApplicationPolicy()
					.getReplication();
			Map<String, Short> dcBrokerCount = new HashMap<String, Short>();
			for (String broker : aliveBrokers) {
				String dc = ApplicationConstants.HOST_DC_MAPPING.get(broker);
				if (dcBrokerCount.containsKey(dc)) {
					Short count = dcBrokerCount.get(dc);
					count++;
					dcBrokerCount.put(dc, count);
				} else {
					dcBrokerCount.put(dc, (short) 1);
				}
			}

			if (allBrokers.size() < numBrokersRequiredPerDC
					* messageBroker.getApplicationPolicy()
							.getDataCentersHosted().size()) {
				logger.info("Not all brokers have been started initially!");
			}

			for (String hostedDc : messageBroker.getApplicationPolicy()
					.getDataCentersHosted()) {
				Short currentNumberOfBrokers = 0;
				if (dcBrokerCount.containsKey(hostedDc)) {
					currentNumberOfBrokers = dcBrokerCount.get(hostedDc);
				}
				logger.info("DC :" + hostedDc + " Count: "
						+ currentNumberOfBrokers);

				if (dcScaledTimeMap.containsKey(hostedDc)) {
					Date lastScaledTime = dcScaledTimeMap.get(hostedDc);
					if (!CommonUtil.canInitiateScaling(lastScaledTime,
							currentTime)) {
						continue;
					}
				}

				int instancesToScale = numBrokersRequiredPerDC
						- currentNumberOfBrokers;
				if (instancesToScale > 0) {
					logger.info("Need to scale " + instancesToScale
							+ " hosts in " + hostedDc);
					CloudServiceClient.getCloudServiceClient()
							.scaleBrokerInstances(hostedDc, instancesToScale);
					dcScaledTimeMap.put(hostedDc, currentTime);
				}
			}

		} catch (KeeperException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (InterruptedException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}

	}

}
